from langchain_core.runnables import RunnableWithMessageHistory, RunnablePassthrough
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.chat_message_histories import ChatMessageHistory
import os

# 定义模型
llm = ChatOpenAI(model_name="deepseek-chat", api_key=os.environ["DEEPSEEK_API_KEY"],base_url=os.environ["DEEPSEEK_BASE_URL"])
# 定义提示词模板
prompt_template = ChatPromptTemplate.from_messages([
    ('system','你是一个乐于助人的助手，请尽你所能回答问题'),
    MessagesPlaceholder(variable_name="history_msg"), # 获取用户的聊天记录  variable_name 为 RunnableWithMessageHistory中的input_messages_key
    ('human',"{input_msg}")
])

temp_chat_history =  ChatMessageHistory()
temp_chat_history.add_user_message("你好，我叫bruce")
temp_chat_history.add_ai_message("你好 bruce，我是deepseek-chat")
temp_chat_history.add_user_message("我今天心情很不错")
temp_chat_history.add_ai_message("今天，你的心情如何")
temp_chat_history.add_user_message("今天早上，我做了几个卡布奇诺")
temp_chat_history.add_ai_message("你上午做了什么食物")

chain = prompt_template | llm

chain_with_message_history =RunnableWithMessageHistory(
    chain,
    lambda session_id:temp_chat_history,
    input_messages_key="input_msg", # 自定义输入的聊天记录的key
    history_messages_key="history_msg" # 对应的MessagesPlaceholder中的variable_name 历史聊天记录
)

# 将聊天记录浓缩成大模型提示词
def summarize_messages(chain_input):
    stored_messages = temp_chat_history.messages
    if len(stored_messages) == 0:
        return False
    summarization_prompt = ChatPromptTemplate.from_messages([
        ('user','请将上述聊天记录浓缩成一条摘要消息，尽可能包含更多具体细节'),
        MessagesPlaceholder(variable_name="summarize_history_msg"),
    ])

    summarization_chain = summarization_prompt | llm
    summary_message = summarization_chain.invoke({"summarize_history_msg":stored_messages})
    temp_chat_history.clear()
    temp_chat_history.add_message(summary_message)
    return True

# 触发 总结聊天记录的操作
chain_with_summarization = (RunnablePassthrough.assign(messages_trimmed=summarize_messages)
                           | chain_with_message_history)

config = {'configurable':{'session_id':'boss'}}
response = chain_with_summarization.invoke({"input_msg":"名字,下午在干啥，心情"}, config)

print("response:",response)
print("temp_chat_history:",temp_chat_history.messages)